#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/rand.h>
#include <openssl/err.h>

#define AES_KEY_SIZE 256
#define AES_IV_SIZE 16
#define BUFFER_SIZE 2048

void print_openssl_error(const char *msg) {
    fprintf(stderr, "❌ %s\n", msg);
    ERR_print_errors_fp(stderr);
}

void generate_rsa_keypair(const char *priv_file, const char *pub_file, const char *who) {
    if (access(priv_file, F_OK) == 0 && access(pub_file, F_OK) == 0) {
        printf("🔑 %s 密钥已存在，跳过生成。\n", who);
        return;
    }

    printf("🔧 生成 %s 的 RSA 密钥对...\n", who);
    EVP_PKEY *pkey = EVP_PKEY_new();
    RSA *rsa = RSA_new();
    BIGNUM *bn = BN_new();
    BN_set_word(bn, RSA_F4);
    RSA_generate_key_ex(rsa, 2048, bn, NULL);
    EVP_PKEY_assign_RSA(pkey, rsa);

    FILE *priv = fopen(priv_file, "wb");
    PEM_write_PrivateKey(priv, pkey, NULL, NULL, 0, NULL, NULL);
    fclose(priv);

    FILE *pub = fopen(pub_file, "wb");
    PEM_write_PUBKEY(pub, pkey);
    fclose(pub);

    EVP_PKEY_free(pkey);
    BN_free(bn);

    printf("✅ %s 密钥生成完成\n", who);
}

int main() {
    ERR_load_crypto_strings();

    // Alice 自己的密钥
    generate_rsa_keypair("alice_private.key", "alice_public.key", "Alice");

    // 检查 Bob 公钥是否存在
    if (access("bob_public.key", F_OK) != 0) {
        printf("⚠️ 未检测到 Bob 的公钥 (bob_public.key)。请先运行 ./bob 生成密钥。\n");
        return 1;
    }

    // 明文
    unsigned char plaintext[] = "这里是20231325jlq（Alice）发送给20231311wgq（Bob）的一条加密消息。";
    printf("\n📨 原始明文: %s\n", plaintext);

    // 生成 AES 密钥与 IV
    unsigned char aes_key[AES_KEY_SIZE / 8];
    unsigned char aes_iv[AES_IV_SIZE];
    RAND_bytes(aes_key, sizeof(aes_key));
    RAND_bytes(aes_iv, sizeof(aes_iv));

    // AES 加密
    unsigned char ciphertext[BUFFER_SIZE];
    int len, ciphertext_len;
    EVP_CIPHER_CTX *aes_ctx = EVP_CIPHER_CTX_new();
    EVP_EncryptInit_ex(aes_ctx, EVP_aes_256_cbc(), NULL, aes_key, aes_iv);
    EVP_EncryptUpdate(aes_ctx, ciphertext, &len, plaintext, strlen((char *)plaintext));
    ciphertext_len = len;
    EVP_EncryptFinal_ex(aes_ctx, ciphertext + len, &len);
    ciphertext_len += len;
    EVP_CIPHER_CTX_free(aes_ctx);

    // 加载 Bob 公钥
    FILE *bob_pub_file = fopen("bob_public.key", "r");
    EVP_PKEY *bob_pubkey = PEM_read_PUBKEY(bob_pub_file, NULL, NULL, NULL);
    fclose(bob_pub_file);

    // 用 Bob 公钥加密 AES 密钥
    EVP_PKEY_CTX *enc_ctx = EVP_PKEY_CTX_new(bob_pubkey, NULL);
    EVP_PKEY_encrypt_init(enc_ctx);
    EVP_PKEY_CTX_set_rsa_padding(enc_ctx, RSA_PKCS1_PADDING);

    unsigned char envelope[BUFFER_SIZE];
    size_t envelope_len = sizeof(envelope);
    EVP_PKEY_encrypt(enc_ctx, envelope, &envelope_len, aes_key, sizeof(aes_key));
    EVP_PKEY_CTX_free(enc_ctx);
    EVP_PKEY_free(bob_pubkey);

    // Alice 用私钥签名
    FILE *alice_priv_file = fopen("alice_private.key", "r");
    EVP_PKEY *alice_privkey = PEM_read_PrivateKey(alice_priv_file, NULL, NULL, NULL);
    fclose(alice_priv_file);

    EVP_MD_CTX *sign_ctx = EVP_MD_CTX_new();
    unsigned char signature[BUFFER_SIZE];
    size_t signature_len = sizeof(signature);
    EVP_DigestSignInit(sign_ctx, NULL, EVP_sha256(), NULL, alice_privkey);
    EVP_DigestSignUpdate(sign_ctx, plaintext, strlen((char *)plaintext));
    EVP_DigestSignFinal(sign_ctx, signature, &signature_len);
    EVP_MD_CTX_free(sign_ctx);
    EVP_PKEY_free(alice_privkey);

    // 写出文件
    FILE *out = fopen("message_to_bob", "wb");
    int iv_len = AES_IV_SIZE;
    fwrite(&iv_len, sizeof(int), 1, out);
    fwrite(aes_iv, 1, AES_IV_SIZE, out);
    fwrite(&ciphertext_len, sizeof(int), 1, out);
    fwrite(ciphertext, 1, ciphertext_len, out);
    fwrite(&envelope_len, sizeof(size_t), 1, out);
    fwrite(envelope, 1, envelope_len, out);
    fwrite(&signature_len, sizeof(size_t), 1, out);
    fwrite(signature, 1, signature_len, out);
    fclose(out);

    printf("\n✅ 数字信封已生成: message_to_bob\n");
    ERR_free_strings();
    return 0;
}
